# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: inkyu

import logging

from coin.base.book import BookBuilderBase
from coin.base.book.aggregated_book_builder import BookUpdateSide
from coin.exchange.okex.kr_rest.product import OkexProduct


class BookEntry(object):
  def __init__(self, price, qty):
    self.price = float(price)
    self.qty = float(qty)


class BookBuilder(BookBuilderBase):
  ProductType = OkexProduct

  def __init__(self, native_symbol, current_datetime=None, book_builder_observer=None):
    super().__init__("Okex",
                     native_symbol,
                     current_datetime,
                     book_builder_observer=book_builder_observer)

  def update_fullbook(self, bids: list, asks: list):
    self._ask_price_map.clear()
    self._bid_price_map.clear()

    for level in asks:
      price = float(level[0])
      qty = float(level[1])
      if qty > 0:
        self._insert_ask(price, qty, True)

    for level in bids:
      price = float(level[0])
      qty = float(level[1])
      if qty > 0:
        self._insert_bid(price, qty, True)

  def get_diffs(self, bids: list, asks: list):
    diffs = []
    for level in bids:
      price = float(level[0])
      qty = float(level[1])
      bqty = self._bid_price_map.get(price, BookEntry(price, 0)).qty
      diffs.append([BookUpdateSide.BID, price, qty - bqty])

    for level in asks:
      price = float(level[0])
      qty = float(level[1])
      bqty = self._ask_price_map.get(price, BookEntry(price, 0)).qty
      diffs.append([BookUpdateSide.ASK, price, qty - bqty])
    return diffs

  def update_diff(self, bids: list, asks: list):
    insert_max_bid_price = None
    insert_min_ask_price = None

    for level in asks:
      price = float(level[0])
      qty = float(level[1])
      if qty > 0:
        self._insert_ask(price, qty, False)
        if insert_min_ask_price is None or price < insert_min_ask_price:
          insert_min_ask_price = price
      else:
        self._delete_ask(price)

    for level in bids:
      price = float(level[0])
      qty = float(level[1])
      if qty > 0:
        self._insert_bid(price, qty, False)
        if insert_max_bid_price is None or price > insert_max_bid_price:
          insert_max_bid_price = price
      else:
        self._delete_bid(price)

    diffs_add = []
    if (insert_max_bid_price is not None and self.ask0() is not None
        and insert_max_bid_price >= self.ask0().price):
      diffs_add.extend(self._delete_ask_up_to_and_including(insert_max_bid_price))
    if (insert_min_ask_price is not None and self.bid0() is not None
        and insert_min_ask_price <= self.bid0().price):
      diffs_add.extend(self._delete_bid_up_to_and_including(insert_min_ask_price))
    return diffs_add

  def _insert_bid(self, price, qty, insert_only=False):
    assert qty > 0
    entry = BookEntry(price, qty)
    if insert_only and price in self._bid_price_map:
      logging.error('Bid book already has the price: %.f', price)
    self._bid_price_map[price] = entry

  def _insert_ask(self, price, qty, insert_only=False):
    assert qty > 0
    entry = BookEntry(price, qty)
    if insert_only and price in self._ask_price_map:
      logging.error('Ask book already has the price: %.f', price)
    self._ask_price_map[price] = entry

  def _delete_bid(self, price):
    diffs = []
    if price in self._bid_price_map:
      diffs.append((BookUpdateSide.BID, price, -self._bid_price_map[price].qty))
      del self._bid_price_map[price]
    return diffs

  def _delete_ask(self, price):
    diffs = []
    if price in self._ask_price_map:
      diffs.append((BookUpdateSide.ASK, price, -self._ask_price_map[price].qty))
      del self._ask_price_map[price]
    return diffs

  def _delete_bid_up_to_and_including(self, price):
    elist = self._bid_price_map.get_entries(price, float('Inf'), exclude_max=False)
    plist = [p.price for p in elist]
    diffs = []
    for p in plist:
      diffs.extend(self._delete_bid(p))
    return diffs

  def _delete_ask_up_to_and_including(self, price):
    elist = self._ask_price_map.get_entries(-float('Inf'), price, exclude_max=False)
    plist = [p.price for p in elist]
    diffs = []
    for p in plist:
      diffs.extend(self._delete_ask(p))
    return diffs
